package main

import (
	"fmt"
	"log"
	"os"
	"os/signal"

	"github.com/Shopify/sarama"

	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/linkedin/goavro"
)

func main() {
	codec, err := goavro.NewCodec(
		`{"type":"record","name":"networkData","namespace":"com.lee","fields":[{"name":"NetworkAccessLayer","type":{"type":"record","name":"Ethernet","namespace":"com.lee","fields":[{"name":"Length","type":"int"},{"name":"Time","type":"long"},{"name":"SrcMAC","type":"string"},{"name":"DstMAC","type":"string"},{"name":"NetworkProtocol","type":"string"}]}},{"name":"InternetLayer","type":["null",{"type":"record","name":"IP","namespace":"com.lee","fields":[{"name":"IHL","type":"int"},{"name":"TOS","type":"int"},{"name":"Id","type":"int"},{"name":"Flags","type":"string"},{"name":"FragOffset","type":"int"},{"name":"TTL","type":"int"},{"name":"SrcIP","type":"string"},{"name":"DstIP","type":"string"}]},{"type":"record","name":"ICMP","namespace":"com.lee","fields":[{"name":"Id","type":"int"},{"name":"Seq","type":"int"},{"name":"Checksum","type":"int"},{"name":"type","type":"string"}]}]},{"name":"TransportLayer","type":["null",{"type":"record","name":"TCP","namespace":"com.lee","fields":[{"name":"ApplicationProtocol","type":"string"},{"name":"SrcPort","type":"int"},{"name":"DstPort","type":"int"},{"name":"Seq","type":"long"},{"name":"Ack","type":"long"},{"name":"FIN","type":"boolean"},{"name":"SYN","type":"boolean"},{"name":"RST","type":"boolean"},{"name":"PSH","type":"boolean"},{"name":"ACK","type":"boolean"},{"name":"URG","type":"boolean"},{"name":"ECE","type":"boolean"},{"name":"CWR","type":"boolean"},{"name":"NS","type":"boolean"},{"name":"Window","type":"int"},{"name":"CheckSum","type":"int"}]},{"type":"record","name":"UDP","namespace":"com.lee","fields":[{"name":"SrcPort","type":"int"},{"name":"DstPort","type":"int"},{"name":"Length","type":"int"},{"name":"CheckSum","type":"int"},{"name":"ApplicationProtocol","type":"string"}]}]}]}`)
	if err != nil {
		panic(err)
	}
	path := "/Users/liruopeng/Downloads/test2.pcap"

	// Trap SIGINT to trigger a shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)

	var (
		successes, errors int
	)

	producer := initProducer("10.36.8.128:9092")

	go func() {
		for range (*producer).Successes() {
			fmt.Println(successes)
			successes++
		}
	}()

	go func() {
		for err := range (*producer).Errors() {
			log.Println(err)
			errors++
		}
	}()

	sendToKafka(path, producer, codec)

WaitLoop:
	for {
		select {
		case <-signals:
			(*producer).AsyncClose()
			break WaitLoop
		}
	}
	fmt.Println(successes, errors)
}

func sendToKafka(path string, producer *sarama.AsyncProducer, codec *goavro.Codec) {
	handle, err := pcap.OpenOffline(path)
	if err != nil {
		log.Fatal(err)
	}
	defer handle.Close()
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())

	for packet := range packetSource.Packets() {
		packetMap := GetAvroMap(packet)
		// log.Println(packetMap)
		// if len(packetMap) < 29 {
		// 	continue
		// }
		binary, err := codec.BinaryFromNative(nil, packetMap)
		if err != nil {
			panic(err)
		}
		fmt.Println(binary)
		message := &sarama.ProducerMessage{Topic: "multirecord3", Value: sarama.ByteEncoder(binary)}

		(*producer).Input() <- message

	}
}

func initProducer(url string) *sarama.AsyncProducer {
	config := sarama.NewConfig()
	config.Producer.MaxMessageBytes = 1000000
	// config.Producer.Return.Successes = true
	producer, err := sarama.NewAsyncProducer([]string{url}, config)
	if err != nil {
		panic(err)
	}
	return &producer
}

func GetAvroMap(packet gopacket.Packet) map[string]interface{} {

	data := map[string]interface{}{}
	ethernetData := GetEthernet(&packet)
	data["NetworkAccessLayer"] = ethernetData

	networkProtocol := ethernetData["NetworkProtocol"]
	internetLayerData := map[string]interface{}{}
	switch networkProtocol {
	case "IPv4":
		{
			internetLayerData = GetIpData(&packet)
			data["InternetLayer"] = goavro.Union("com.lee.IP", internetLayerData)

		}
	case "ICMPv4":
		{
			data["InternetLayer"] = goavro.Union("com.lee.ICMP", GetICMPData(&packet))
		}
	default:
		{
			data["InternetLayer"] = goavro.Union("null", nil)
		}
	}

	transportLayerProtocol := internetLayerData["TransportProtol"]
	switch transportLayerProtocol {
	case "TCP":
		{
			data["TransportLayer"] = goavro.Union("com.lee.TCP", GetTCPData(&packet))
		}
	case "UDP":
		{
			data["TransportLayer"] = goavro.Union("com.lee.UDP", GetUDPData(&packet))
		}
	default:
		{
			data["TransportLayer"] = goavro.Union("null", nil)
		}
	}

	return data
}

func GetEthernet(packet *gopacket.Packet) map[string]interface{} {
	ethernet := map[string]interface{}{}
	ethernet["Length"] = len((*packet).Data())
	ethernet["Time"] = (*packet).Metadata().Timestamp.Unix()

	// 网络接口层
	ethernetLayer := (*packet).Layer(layers.LayerTypeEthernet)
	ethernetPacket, _ := ethernetLayer.(*layers.Ethernet)

	ethernet["SrcMAC"] = ethernetPacket.SrcMAC.String()
	ethernet["DstMAC"] = ethernetPacket.DstMAC.String()
	ethernet["NetworkProtocol"] = ethernetPacket.EthernetType.String()

	return ethernet
}

func GetIpData(packet *gopacket.Packet) map[string]interface{} {
	ipData := map[string]interface{}{}
	// 网络层

	ipLayer := (*packet).Layer(layers.LayerTypeIPv4)
	ip, _ := ipLayer.(*layers.IPv4)

	ipData["IHL"] = int(ip.IHL)
	ipData["TOS"] = int(ip.TOS)
	ipData["Id"] = int(ip.Id)
	ipData["Flags"] = ip.Flags.String()
	ipData["FragOffset"] = int(ip.FragOffset)
	ipData["TTL"] = int(ip.TTL)
	ipData["SrcIP"] = ip.SrcIP.String()
	ipData["DstIP"] = ip.DstIP.String()
	ipData["TransportProtol"] = ip.NextLayerType().String()

	return ipData
}

func GetTCPData(packet *gopacket.Packet) map[string]interface{} {
	// 传输层 TCP
	tcpData := map[string]interface{}{}
	tcpLayer := (*packet).Layer(layers.LayerTypeTCP)
	tcp, _ := tcpLayer.(*layers.TCP)
	tcpData["SrcPort"] = int(tcp.SrcPort)
	tcpData["DstPort"] = int(tcp.DstPort)
	tcpData["ApplicationProtocol"] = "unknown"
	if name, ok := layers.TCPPortNames[tcp.SrcPort]; ok {
		tcpData["ApplicationProtocol"] = name
	}

	if name, ok := layers.TCPPortNames[tcp.DstPort]; ok {
		tcpData["ApplicationProtocol"] = name
	}
	tcpData["Seq"] = int64(tcp.Seq)
	tcpData["Ack"] = int64(tcp.Ack)
	tcpData["FIN"] = tcp.FIN
	tcpData["SYN"] = tcp.SYN
	tcpData["RST"] = tcp.RST
	tcpData["PSH"] = tcp.PSH
	tcpData["ACK"] = tcp.ACK
	tcpData["URG"] = tcp.URG
	tcpData["ECE"] = tcp.ECE
	tcpData["CWR"] = tcp.CWR
	tcpData["NS"] = tcp.NS
	tcpData["Window"] = int(tcp.Window)
	tcpData["CheckSum"] = int(tcp.Checksum)
	return tcpData
}

func GetUDPData(packet *gopacket.Packet) map[string]interface{} {
	// 传���层 TCP
	udpData := map[string]interface{}{}
	udpLayer := (*packet).Layer(layers.LayerTypeUDP)
	udp, _ := udpLayer.(*layers.UDP)
	udpData["SrcPort"] = int(udp.SrcPort)
	udpData["DstPort"] = int(udp.DstPort)
	udpData["ApplicationProtocol"] = "unknown"
	udpData["Length"] = int(udp.Length)
	if name, ok := layers.UDPLitePortNames[layers.UDPLitePort(udp.SrcPort)]; ok {
		udpData["ApplicationProtocol"] = name
	}

	if name, ok := layers.UDPLitePortNames[layers.UDPLitePort(udp.DstPort)]; ok {
		udpData["ApplicationProtocol"] = name
	}
	udpData["CheckSum"] = int(udp.Checksum)
	return udpData
}

func GetICMPData(packet *gopacket.Packet) map[string]interface{} {
	// 传输层 TCP
	icmpData := map[string]interface{}{}
	icmpLayer := (*packet).Layer(layers.LayerTypeICMPv4)
	icmp, _ := icmpLayer.(*layers.ICMPv4)

	icmpData["Id"] = int(icmp.Checksum)
	icmpData["Seq"] = int(icmp.Checksum)
	icmpData["CheckSum"] = int(icmp.Checksum)

	icmpData["type"] = icmp.TypeCode.String()
	return icmpData
}
